import transaction
from stockforecaster.lib import ystockquote
from stockforecaster.lib.utils import split_cslist
from stockforecaster.models import Company
from stockforecaster.tasks import fetch_stock_quotes
from wtforms import Form, validators
from wtforms.fields import TextAreaField


def filter_symbol(symbol):
    if symbol:
        return symbol.upper()
    

class AddCompaniesForm(Form):
    
    symbol = TextAreaField('Stock ticker symbols', [
        validators.Required(message='Must enter at least one ticker symbol')],
        filters=[filter_symbol])

    def validate_symbol(self, field):
        for symbol in split_cslist(field.data):
            if not ystockquote.symbol_exists(symbol):
                raise validators.ValidationError('Invalid ticker symbol %s' % symbol)
        
    def add_companies(self):
        companies_added = 0
        symbols = split_cslist(self.symbol.data)
        for symbol in symbols:
            if Company.get_by_symbol(symbol):
                continue
            name = ystockquote.get_company_name(symbol)
            exchange = ystockquote.get_stock_exchange(symbol)
            company = Company.create(symbol=symbol, name=name, exchange=exchange)
            fetch_stock_quotes.apply_async(args=[company.id], countdown=3*len(symbols))
            companies_added += 1
        return companies_added
